from dotenv import load_dotenv
from langchain.chains import GraphCypherQAChain
from langchain_community.graphs import Neo4jGraph
from langchain.prompts import ChatPromptTemplate
from model.my_chat_model import ChatModel
from langchain_core._api import LangChainDeprecationWarning
import warnings
import os

warnings.filterwarnings("ignore", category=LangChainDeprecationWarning)

class RAG_Neo4jService:
    def __init__(self):
        #加载环境变量
        load_dotenv()
        #初始化聊天模型
        chat=ChatModel()
        self.llm=chat.get_line_model()
        #初始化聊天数据库
        self.graph=Neo4jGraph(
            url=os.getenv("NEO4J_URI"),
            username=os.getenv("NEO4J_USERNAME"),
            password=os.getenv("NEO4J_PASSWORD"),
        )
    #创建一个对话链
    def create_chain(self):
        #获取提示模板
        prompt = ChatPromptTemplate.from_messages([
            ("system", """
                你是一个专业的 Cypher 查询生成器，能够根据用户的自然语言问题生成正确的 Cypher 查询语句。
                请严格遵循以下数据库模式：

                节点类型与属性：
                - Company: reg_date, risk_level, reg_no, name, industry, legal_rep, registered_capital, status
                - Senior: id_no, gender, birth_date, name, position
                - Shareholder: id_no, nationality, is_controlled, name, type

                关系类型与属性：
                - (Shareholder)-[:HOLDS]->(Company)：该关系有一个属性 `ratio`（浮点数，表示持股比例）
                - (Company)-[:INVESTS_IN]->(Company)：该关系有一个属性 `amount`（整数或浮点数，表示投资金额）
                - (Shareholder)-[:RELATED_TO]->(Shareholder)：该关系有一个属性 `relation`（字符串，如"亲属"、"合作伙伴"等）
                - (Senior)-[:SERVES_AS]->(Company)：该关系有一个属性 `role`（字符串，表示职务）

                生成 Cypher 查询时请注意：
                1. 必须使用 MATCH 子句查找对应节点和关系
                2. 使用 RETURN 返回所需字段（如 h.ratio, i.amount, r.relation, a.role）
                3. 不要返回整个节点，只返回需要的属性
                4. 所有字符串比较使用双引号包围
                5. 示例：
                   - 查询张三持股比例：MATCH (s:Shareholder {{name: "张三"}})-[h:HOLDS]->(c:Company {{name: "某公司"}}) RETURN h.ratio
                   - 查询李四担任职位：MATCH (s:Senior {{name: "李四"}})-[a:SERVES_AS]->(c:Company) RETURN a.role, c.name
                   - 查询股东A和股东B之间的关系:MATCH (s1:Shareholder {{name: "股东A"}})-[r:RELATED_TO ]->(s2:Shareholder {{name: "股东B"}}) RETURN r.relation
                   - 查询公司A在公司B投资了多少资金：MATCH (c1:Company {{name: "公司A"}})-[i:INVESTS_IN]->(c:Company {{name:"公司B"}}) RETURN i.amount
                """),
            ("human", "{question}")
        ])
        #创建链
        chain=GraphCypherQAChain.from_llm(
            llm=self.llm,
            graph=self.graph,
            allow_dangerous_requests=True,
            cypher_prompt=prompt,
            verbose=True,
            top_k=20,
            validate_cypher=True
        )
        return chain